import os
import sys
# 添加父级目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.insert(0, parent_dir)
from FlagEmbedding import FlagModel
from my_common import BGE_LARGE_ZH_V1_5_MODEL

os.environ['TRANSFORMERS_NO_ADVISORY_WARNINGS'] = 'true'
# single GPU is better for small tasks
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

sentences = [
    "That is a happy dog",
    "That is a very happy person",
    "Today is a sunny day",
]


if __name__ == '__main__':
    # get the BGE embedding model
    model = FlagModel(BGE_LARGE_ZH_V1_5_MODEL)
    # encode the queries and corpus
    embeddings = model.encode(sentences)
    print(f"Embeddings:\n{embeddings.shape}")

    scores = embeddings @ embeddings.T
    print(f"Similarity scores:\n{scores}")